import pandas as pd
import numpy as np
import math
import matplotlib.pyplot as plt
from pythonds.basic import Stack
from scipy.spatial.kdtree import KDTree

def readDataSet(filename, div=2):
    data = np.loadtxt(filename)
    return data

def readDataSetwithLabel(filename, div=2):
    matrix = np.loadtxt(filename)
    label = matrix[:, -1]
    matrix = np.delete(matrix, -1, axis=1)
    return matrix, label

def readLabel(filename):
    data = pd.read_csv(filename, header=None)
    matrix = np.array(data)
    return matrix

def getKNN(matrix, k):
    '''
    获得knn连接图
    :param matrix:  坐标信息
    :param k: k个最近邻
    :return: dd: 距离矩阵，最近的k个点的距离
    :return: ii: 索引矩阵，最近的k个点的索引
    '''
    kd_tree = KDTree(matrix)
    dd, ii = kd_tree.query(matrix, k = k+1)
    dd = np.delete(dd, 0, axis=1)
    ii = np.delete(ii, 0, axis=1)
    return dd, ii

def getConnectedGraph(dd, length, k):
    '''
    通过距离矩阵喝索引矩阵获得连通图并计算对应的权值，
    由于计算权值时使用除以最近邻距离，所以连接时不连接最近邻
    :param dd: k近邻距离矩阵
    :param length:
    :param k:
    :return: w: 权值矩阵
    '''
    w = np.zeros((length, int(k)), dtype=float)
    for i in range(length):
        for j in range(k):
            w[i][j] = -1.0 * dd[i][j]
    return w

def getPointW(w, length):
    '''
    将权重矩阵计算平均值，代表每个点的权重
    :param w:
    :param length:
    :return:
    '''
    weight = np.mean(w, axis=1).reshape((length, 1))
    return weight

def getProbwitInterval(obj, length, p):
    '''
    使用插值法表示obj的数值分布
    :param obj: 待表示的一维数组
    :param length: 数组长度
    :param p: 间隔
    :return:
    '''
    number = np.zeros((p + 1, 1), dtype=float)
    w_max = np.max(obj)
    w_min = np.min(obj)
    dc = (w_max - w_min) / p
    for i in range(length):
        j = math.floor(float((obj[i] - w_min) / dc))
        number[j] = number[j] + 1
    prob = number / length
    return prob

def pltDecisionGraph(obj, prob, filename=None):
    '''
    绘制决策图，并将图片存在filename
    :param obj:
    :param prob:
    :param filename:
    :return:
    '''
    p = len(prob) - 1
    w_max = np.max(obj)
    w_min = np.min(obj)
    dc = (w_max - w_min) / p

    fig = plt.figure()
    x_ = np.arange(0, p + 1) * dc + w_min
    plt.plot(x_, prob, color='r')
    if not filename:
        plt.savefig(filename)
    plt.show()

def cutOffNosie(w_matrix, w_):
    '''
    根据第一张决策图切分连通图
    :param w_matrix: 权值矩阵
    :param w_: 阈值
    :return:
    '''
    l = w_matrix.shape[0]
    k = w_matrix.shape[1]

    for i in range(l):
        for j in range(k):
            if w_matrix[i][j] < w_:
                w_matrix[i][j] = 0
    return

def clustering(w, posi, k):
    # 使用广度深度搜索的方法寻找连通子图，每一个子图是一个类
    length = w.shape[0]
    visited = np.zeros((length, 1), dtype=int)
    s = Stack()
    cluster = np.zeros((length, 1), dtype=int)
    count = 1
    for i in range(length):
        if visited[i] == 0:
            visited[i] = 1
            if w[i][0] == 0:
                #如果该点没有任何连接点，认为是一个nosie point
                cluster[i] = -1
                continue
            else:
                s.push(i)
                while not s.isEmpty():
                    top = s.pop()
                    cluster[top] = count
                    for j in range(k-1):
                        if w[top][j] != 0:
                            if visited[posi[top][j]] == 0:
                                visited[posi[top][j]] = 1
                                s.push(posi[top][j])
                count = count + 1
    return cluster

from sklearn.decomposition import PCA

def plotResult(matrix, result, filename, k, w_):
    pca = PCA(n_components=2)  # n_components设置降维后的特征数
    pca.fit(matrix)  # 拟合
    location = pca.transform(matrix)  # 获取新矩阵

    length = len(result)
    markers = ['.', '*', '+', 'x', '^']
    colors = ['maroon', 'red', 'peru', 'gold', 'olive', 'yellowgreen', 'lawngreen', 'springgreen']
    colors = colors + ['turquoise', 'teal', 'deepskyblue', 'dodgerblue', 'royalblue', 'navy']
    colors = colors + ['slategrey', 'orchid', 'm', 'deeppink', 'crimson']
    plt.figure(figsize=(10, 10))
    plt.title("TGC k={%d} w*={:-.2f}".format(k, w_))
    for i in range(0, length):
        index = int(result[i])
        if index == -1:
            plt.plot(location[i][0], location[i][1], color=(0,0,0),  marker='.')
        else:
            plt.plot(location[i][0], location[i][1], color=colors[index%19], marker=markers[index%5])
    plt.xlabel('Attribute 1'), plt.ylabel('Attribute 2')
    plt.savefig(filename)
    plt.show()

def tpg(input_file, output_directory, k, file_withlabel):
    k = k
    filename = input_file

    if file_withlabel:
        data, label = readDataSetwithLabel(filename)
    else:
        data = readDataSet(filename)

    length = data.shape[0]

    dd, ii = getKNN(data, k)
    w_matrix = getConnectedGraph(dd, length, k)
    w = getPointW(w_matrix, length)

    p_1 = int(length / 10)
    decision_graph = output_directory + '/decision_graph_k=' + str(k) + '.png'
    prob_1 = getProbwitInterval(w, length, p_1)
    pltDecisionGraph(w, prob_1, decision_graph)

    w_ = float(input("input w_\n"))
    cutOffNosie(w_matrix, w_)

    result = clustering(w_matrix, ii, k)
    output = output_directory + '/result_k=' + str(k) + '_w=' + str(w_) + 'MAX.png'
    plotResult(data, result, output, k, w_)
    return result

from getRI import rand_index
from getNMI import printNMI

if __name__ == '__main__':
    k = 10
    filename = './'
    directory = './'
    file_with_label = 1

    result = tpg(filename, directory, k, file_with_label)

    output_file = "tgc_result.txt"
    np.savetxt(output_file, result)

    nmi = printNMI("label.txt", output_file)
    ri = rand_index("label.txt", output_file)

    np.savetxt("k={}__nmi={:.2f}_ri={:.2f}_".format(int(k), nmi, ri) + output_file, result)


